#!/usr/bin/env python
# coding: utf-8

import os
from import_helper import *
import matplotlib.pyplot as plt
import numpy as np
import statsmodels.api as sm
import statsmodels.formula.api as smf
import scipy 


rcParams['figure.figsize'] = (13.0, 7.0)
rcParams['lines.linewidth'] = 3
rcParams['font.size'] = 18
treatment_list_orig = treatment_list


def pretty_plot(ax, treatments=treatment_list):
    n = len(ax.lines)
    list_style = ['-', '--', ':', '-.']
    for i, s in enumerate(list_style[:n]):
        ax.lines[n-1-i].set_linestyle(s)
        ax.lines[n-1-i].set_label(treatments[i])
    plt.legend(ax.lines[::-1], treatments[:n], title='treatment')


# # 7 Sessions q=0 - Random, Priority + No Info, Priority + Info


treatment_list = [treatment_list_orig[0],treatment_list_orig[1],treatment_list_orig[3]]
df = df_mturk_noq.copy()
df_woke = filter_data(df, woke=True, experience=is_greater_than(0))
df = filter_data(df, experience=is_greater_than(0))


# # Table OE1


df_woke['targeted_info'] = (df_woke[CB.treatment] == treatment_list[2])
df_woke['priority_only'] = (df_woke[CB.treatment] == treatment_list[1])
df_woke['random'] = (df_woke[CB.treatment] == treatment_list[0])
df_woke['cluster'] = df_woke[[CB.session, CB.treatment]].apply(
    lambda x: str(x[0]) + '|' + x[1], axis=1)
df_woke_forreg = pd.get_dummies(df_woke, columns=[CB.session], drop_first=True)



formula = 'settled ~ 1 + priority_only + targeted_info'
print(formula)
this_data = df_woke_forreg
model = OLS.from_formula(
    formula,
    data=this_data
)
res = model.fit(cov_type='clustered', clusters=this_data['cluster'])
print(res.summary.as_text())

# Print Table 6
with open("figs/tableOE1_col1.txt", "w") as text_file:
    print(res.summary.as_text(), file=text_file)



df_woke_forreg['accepted_offer'] = df_woke_forreg['player.accepted_offer']
formula0 = 'accepted_offer ~ 1 + priority_only + targeted_info'
this_data = df_woke_forreg
model0 = OLS.from_formula(
    formula0,
    data=this_data
)
res0 = model0.fit(cov_type='clustered', clusters=this_data['cluster'])
print(res0.summary.as_text())
with open("figs/tableOE1_col2.txt", "w") as text_file:
    print(res0.summary.as_text(), file=text_file)


# ## Figure OE2

settled_by_group = df_woke.groupby(unit_playergroup_rnd)[[CB.settled, CB.num_treatment]].mean().reset_index()
settled_by_group[CB.treatment] = settled_by_group[CB.num_treatment].apply(to_treatment_name)


ax = sns.ecdfplot(
    settled_by_group, x=CB.settled,
    hue=CB.treatment,
    palette='colorblind',
    hue_order=treatment_list
)

pretty_plot(ax)
plt.legend(title='Treatment', loc='upper left', labels=['Priority+Info',
                                                        'Priority+No-Info',
                                                        'Random'])

plt.xlabel('Settlement Rate in (Session, Treatment, Round)')
plt.ylabel('CDF')
plt.savefig('figs/figureOE2.pdf')


# In[ ]:




